package com.pw.study.flink.watermark;

import com.pw.study.flink.entities.WaterSensor;
import org.apache.flink.api.common.eventtime.Watermark;
import org.apache.flink.api.common.eventtime.WatermarkGenerator;
import org.apache.flink.api.common.eventtime.WatermarkGeneratorSupplier;
import org.apache.flink.api.common.eventtime.WatermarkOutput;
import org.apache.flink.api.common.eventtime.WatermarkStrategy;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.windowing.ProcessWindowFunction;
import org.apache.flink.streaming.api.windowing.assigners.SlidingProcessingTimeWindows;
import org.apache.flink.streaming.api.windowing.time.Time;
import org.apache.flink.streaming.api.windowing.windows.TimeWindow;
import org.apache.flink.util.Collector;

public class WatermarkCustom {

    public static void main(String[] args) {
        wc2();
    }

    private static void wc2() {
        Configuration conf = new Configuration();
        conf.setInteger("rest.port", 2000);
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment(conf);
        env.setParallelism(1);
        //1.数据输出
        SingleOutputStreamOperator<WaterSensor> stream = env.socketTextStream("hadoop111", 9999)
                .map(value -> {
                    String[] data = value.split(",");
                    return new WaterSensor(data[0], Long.valueOf(data[1]), Integer.valueOf(data[2])
                    );
                });
        //2.自定义水位策略
        WatermarkStrategy<WaterSensor> strategy = new WatermarkStrategy<WaterSensor>() {

            public WatermarkGenerator<WaterSensor> createWatermarkGenerator(WatermarkGeneratorSupplier.Context context) {
                return new MyPeriod(3);
            }
        };
        //间歇策略
        WatermarkStrategy<WaterSensor> strategy1 = new WatermarkStrategy<WaterSensor>() {

            public WatermarkGenerator<WaterSensor> createWatermarkGenerator(WatermarkGeneratorSupplier.Context context) {
                return new MyPunctuated(3);
            }
        };


        //3.开始
        stream.assignTimestampsAndWatermarks(strategy).keyBy(WaterSensor::getId)
                .window(SlidingProcessingTimeWindows.of(Time.seconds(5), Time.seconds(5)))
                .process(new ProcessWindowFunction<WaterSensor, String, String, TimeWindow>() {
                    @Override
                    public void process(String key, Context ctx, Iterable<WaterSensor> elements, Collector<String> out) throws Exception {
                        String msg = "当前key: " + key
                                + "窗口: [" + ctx.window().getStart() / 1000 + "," + ctx.window().getEnd() / 1000 + ") 一共有 "
                                + elements.spliterator().estimateSize() + "条数据 ";
                        System.out.println("====== " + msg);
                        out.collect(ctx.window().toString());
                        out.collect(msg);

                    }
                }).print();

        try {
            env.execute();
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    /**
     * 自定义周期性策略
     */
    private static class MyPeriod implements WatermarkGenerator<WaterSensor> {
        //允许最大延迟
        private final long maxDelay;
        private long maxTs = Long.MIN_VALUE;

        public MyPeriod(long maxDelay) {
            this.maxDelay = maxDelay * 1000;
            this.maxTs = Long.MIN_VALUE + this.maxDelay + 1;
        }

        // 每收到一个元素, 执行一次. 用来生产WaterMark中的时间戳
        @Override
        public void onEvent(WaterSensor event, long eventTimestamp, WatermarkOutput output) {
            System.out.println("onEvent...." + eventTimestamp);
            maxTs = Math.max(this.maxTs, eventTimestamp);
            System.out.println("maxTs..." + maxTs);

        }

        //周期性的把WaterMark发射出去, 默认周期是200ms
        @Override
        public void onPeriodicEmit(WatermarkOutput output) {
            // System.out.println("onPeriodicEmit......");
            output.emitWatermark(new Watermark(maxTs - maxDelay - 1));
        }


    }


    //间歇策略
    private static class MyPunctuated implements WatermarkGenerator<WaterSensor> {
        private final Long maxDelay;
        private Long maxTs;

        private MyPunctuated(long maxDelay) {
            this.maxDelay = maxDelay * 1000;
            this.maxTs = Long.MIN_VALUE + this.maxDelay + 1;
        }


        @Override
        public void onEvent(WaterSensor event, long ts, WatermarkOutput output) {
            System.out.println("onEvent...." + ts);
            maxTs = Math.max(this.maxTs, ts);
            output.emitWatermark(new Watermark(maxTs - maxDelay - 1));
        }

        @Override
        public void onPeriodicEmit(WatermarkOutput watermarkOutput) {

        }
    }

}